#include "imu.hpp"
#include "cmsis_os.h"
#include "bsp_dwt.h"
#include "bsp_pwm.h"
#include "bmi088.h"
#include "QuaternionEKF.h"

IMU::IMU(const param_t &param, const calib_t &calib) : param(param), calib(calib) , temperature_pid(temperature_pid_param) {
}

void IMU::Init() {
    IMU_QuaternionEKF_Init(10, 0.001, 10000000, 1, 0);

    BMI088_Init(&hspi1);
    BMI088_SetCalibrateParam(calib.GxOFFSET, calib.GyOFFSET, calib.GzOFFSET, calib.gNORM);
    // BMI088_Calibrate();
}

void IMU::WaitReady() const {
    while (is_ready == false) {
        osDelay(1);
    }
}

void IMU::Update() {
    static uint32_t cnt = 0;
    static uint32_t dwt_cnt = 0;
    static float dt = 0;

    dt = BSP_DWT_GetDeltaT(&dwt_cnt);

    // IMU解算
    if (cnt % 1 == 0) {
        BMI088_Read(&BMI088);

        accel[X] = BMI088.Accel[X];
        accel[Y] = BMI088.Accel[Y];
        accel[Z] = BMI088.Accel[Z];
        gyro[X] = BMI088.Gyro[X];
        gyro[Y] = BMI088.Gyro[Y];
        gyro[Z] = BMI088.Gyro[Z];

        // 用于修正安装误差
        IMU_Param_Correction(&param, gyro, accel);

        // 计算重力加速度矢量和b系的XY两轴的夹角,可用作功能扩展,本demo暂时没用
        atanxz = -atan2f(accel[X], accel[Z]) * 180 / PI;
        atanyz = atan2f(accel[Y], accel[Z]) * 180 / PI;

        // 核心函数,EKF更新四元数
        IMU_QuaternionEKF_Update(gyro[X], gyro[Y], gyro[Z], accel[X],
                                 accel[Y], accel[Z], dt);

        memcpy(q, QEKF_INS.q, sizeof(QEKF_INS.q));

        // 机体系基向量转换到导航坐标系，本例选取惯性系为导航系
        BodyFrameToEarthFrame(xb, xn, q);
        BodyFrameToEarthFrame(yb, yn, q);
        BodyFrameToEarthFrame(zb, zn, q);

        // 将重力从导航坐标系n转换到机体系b，随后根据加速度计数据计算运动加速度
        float gravity_b[3];
        EarthFrameToBodyFrame(gravity, gravity_b, q);
        for (uint8_t i = 0; i < 3; i++) {
            // 过一个低通滤波
            motion_accel_b[i] = (accel[i] - gravity_b[i]) * dt / (accel_lpf + dt) + motion_accel_b[i] * accel_lpf / (
                                    accel_lpf + dt);
        }
        BodyFrameToEarthFrame(motion_accel_b, motion_accel_n, q); // 转换回导航系n

        // 获取最终数据
        yaw = QEKF_INS.Yaw;
        pitch = QEKF_INS.Pitch;
        roll = QEKF_INS.Roll;
        yaw_total_angle = QEKF_INS.YawTotalAngle;
    }

    // IMU温控
    if (cnt % 2 == 0) {
        // 500hz
        IMU_Temperature_Ctrl();
    }

    cnt++;

    is_ready = true;
}

/**
 * @brief reserved.用于修正IMU安装误差与标度因数误差,即陀螺仪轴和云台轴的安装偏移
 *
 *
 * @param param IMU参数
 * @param gyro  角速度
 * @param accel 加速度
 */
void IMU::IMU_Param_Correction(param_t *param, float gyro[3], float accel[3]) {
    static float lastYawOffset, lastPitchOffset, lastRollOffset;
    static float c_11, c_12, c_13, c_21, c_22, c_23, c_31, c_32, c_33;
    float cosPitch, cosYaw, cosRoll, sinPitch, sinYaw, sinRoll;

    if (fabsf(param->yaw - lastYawOffset) > 0.001f ||
        fabsf(param->pitch - lastPitchOffset) > 0.001f ||
        fabsf(param->roll - lastRollOffset) > 0.001f || param->flag) {
        cosYaw = arm_cos_f32(param->yaw / 57.295779513f);
        cosPitch = arm_cos_f32(param->pitch / 57.295779513f);
        cosRoll = arm_cos_f32(param->roll / 57.295779513f);
        sinYaw = arm_sin_f32(param->yaw / 57.295779513f);
        sinPitch = arm_sin_f32(param->pitch / 57.295779513f);
        sinRoll = arm_sin_f32(param->roll / 57.295779513f);

        // 1.yaw(alpha) 2.pitch(beta) 3.roll(gamma)
        c_11 = cosYaw * cosRoll + sinYaw * sinPitch * sinRoll;
        c_12 = cosPitch * sinYaw;
        c_13 = cosYaw * sinRoll - cosRoll * sinYaw * sinPitch;
        c_21 = cosYaw * sinPitch * sinRoll - cosRoll * sinYaw;
        c_22 = cosYaw * cosPitch;
        c_23 = -sinYaw * sinRoll - cosYaw * cosRoll * sinPitch;
        c_31 = -cosPitch * sinRoll;
        c_32 = sinPitch;
        c_33 = cosPitch * cosRoll;
        param->flag = 0;
    }
    float gyro_temp[3];
    for (uint8_t i = 0; i < 3; i++)
        gyro_temp[i] = gyro[i] * param->scale[i];

    gyro[X] = c_11 * gyro_temp[X] +
              c_12 * gyro_temp[Y] +
              c_13 * gyro_temp[Z];
    gyro[Y] = c_21 * gyro_temp[X] +
              c_22 * gyro_temp[Y] +
              c_23 * gyro_temp[Z];
    gyro[Z] = c_31 * gyro_temp[X] +
              c_32 * gyro_temp[Y] +
              c_33 * gyro_temp[Z];

    float accel_temp[3];
    for (uint8_t i = 0; i < 3; i++)
        accel_temp[i] = accel[i];

    accel[X] = c_11 * accel_temp[X] +
               c_12 * accel_temp[Y] +
               c_13 * accel_temp[Z];
    accel[Y] = c_21 * accel_temp[X] +
               c_22 * accel_temp[Y] +
               c_23 * accel_temp[Z];
    accel[Z] = c_31 * accel_temp[X] +
               c_32 * accel_temp[Y] +
               c_33 * accel_temp[Z];

    lastYawOffset = param->yaw;
    lastPitchOffset = param->pitch;
    lastRollOffset = param->roll;
}

/**
 * @brief 温度控制
 *
 */
void IMU::IMU_Temperature_Ctrl() {
    temperature_measure = BMI088.Temperature;
    const float temperature_err = temperature_ref - temperature_measure;
    temperature_pid.CalcPosition(temperature_err);
    BSP_PWM_IMUHeat_SetPower(temperature_pid.out);
}

/**
 * @brief          Transform 3dvector from BodyFrame to EarthFrame
 * @param[1]       vector in BodyFrame
 * @param[2]       vector in EarthFrame
 * @param[3]       quaternion
 */
void IMU::BodyFrameToEarthFrame(const float *vecBF, float *vecEF, const float *q) {
    vecEF[0] = 2.0f * ((0.5f - q[2] * q[2] - q[3] * q[3]) * vecBF[0] +
                       (q[1] * q[2] - q[0] * q[3]) * vecBF[1] +
                       (q[1] * q[3] + q[0] * q[2]) * vecBF[2]);

    vecEF[1] = 2.0f * ((q[1] * q[2] + q[0] * q[3]) * vecBF[0] +
                       (0.5f - q[1] * q[1] - q[3] * q[3]) * vecBF[1] +
                       (q[2] * q[3] - q[0] * q[1]) * vecBF[2]);

    vecEF[2] = 2.0f * ((q[1] * q[3] - q[0] * q[2]) * vecBF[0] +
                       (q[2] * q[3] + q[0] * q[1]) * vecBF[1] +
                       (0.5f - q[1] * q[1] - q[2] * q[2]) * vecBF[2]);
}

/**
 * @brief          Transform 3dvector from EarthFrame to BodyFrame
 * @param[1]       vector in EarthFrame
 * @param[2]       vector in BodyFrame
 * @param[3]       quaternion
 */
void IMU::EarthFrameToBodyFrame(const float *vecEF, float *vecBF, const float *q) {
    vecBF[0] = 2.0f * ((0.5f - q[2] * q[2] - q[3] * q[3]) * vecEF[0] +
                       (q[1] * q[2] + q[0] * q[3]) * vecEF[1] +
                       (q[1] * q[3] - q[0] * q[2]) * vecEF[2]);

    vecBF[1] = 2.0f * ((q[1] * q[2] - q[0] * q[3]) * vecEF[0] +
                       (0.5f - q[1] * q[1] - q[3] * q[3]) * vecEF[1] +
                       (q[2] * q[3] + q[0] * q[1]) * vecEF[2]);

    vecBF[2] = 2.0f * ((q[1] * q[3] + q[0] * q[2]) * vecEF[0] +
                       (q[2] * q[3] - q[0] * q[1]) * vecEF[1] +
                       (0.5f - q[1] * q[1] - q[2] * q[2]) * vecEF[2]);
}

/**
 * @brief        Convert quaternion to eular angle
 */
void IMU::QuaternionToEularAngle(const float *q, float *Yaw, float *Pitch, float *Roll) {
    *Yaw = atan2f(2.0f * (q[0] * q[3] + q[1] * q[2]), 2.0f * (q[0] * q[0] + q[1] * q[1]) - 1.0f) * 57.295779513f;
    *Pitch = atan2f(2.0f * (q[0] * q[1] + q[2] * q[3]), 2.0f * (q[0] * q[0] + q[3] * q[3]) - 1.0f) * 57.295779513f;
    *Roll = asinf(2.0f * (q[0] * q[2] - q[1] * q[3])) * 57.295779513f;
}

/**
 * @brief        Convert eular angle to quaternion
 */
void IMU::EularAngleToQuaternion(float Yaw, float Pitch, float Roll, float *q) {
    float cosPitch, cosYaw, cosRoll, sinPitch, sinYaw, sinRoll;
    Yaw /= 57.295779513f;
    Pitch /= 57.295779513f;
    Roll /= 57.295779513f;
    cosPitch = arm_cos_f32(Pitch / 2);
    cosYaw = arm_cos_f32(Yaw / 2);
    cosRoll = arm_cos_f32(Roll / 2);
    sinPitch = arm_sin_f32(Pitch / 2);
    sinYaw = arm_sin_f32(Yaw / 2);
    sinRoll = arm_sin_f32(Roll / 2);
    q[0] = cosPitch * cosRoll * cosYaw + sinPitch * sinRoll * sinYaw;
    q[1] = sinPitch * cosRoll * cosYaw - cosPitch * sinRoll * sinYaw;
    q[2] = sinPitch * cosRoll * sinYaw + cosPitch * sinRoll * cosYaw;
    q[3] = cosPitch * cosRoll * sinYaw - sinPitch * sinRoll * cosYaw;
}

/**
 * @brief        Update quaternion
 */
void IMU::QuaternionUpdate(float *q, float gx, float gy, float gz, float dt) {
    float qa, qb, qc;

    gx *= 0.5f * dt;
    gy *= 0.5f * dt;
    gz *= 0.5f * dt;
    qa = q[0];
    qb = q[1];
    qc = q[2];
    q[0] += -qb * gx - qc * gy - q[3] * gz;
    q[1] += qa * gx + qc * gz - q[3] * gy;
    q[2] += qa * gy - qb * gz + q[3] * gx;
    q[3] += qa * gz + qb * gy - qc * gx;
}
